/*
   Copyright 2012 Mathias Olsson

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
*/

#include <stdio.h>
#include <stdlib.h>
#include <assert.h>
#include <string.h>

#include <event2/event.h>
#include <event2/listener.h>
#include <event2/buffer.h>
#include <event2/bufferevent.h>
#include <event2/util.h>
#include <event2/dns.h>

#include <polarssl/ssl.h>
#include <polarssl/entropy.h>
#include <polarssl/ctr_drbg.h>

#include <lua.h>
#include <lualib.h>
#include <lauxlib.h>

// Generated by SCons
#include "config.h"

#include "listener.h"
#include "cleartext_stream.h"
#include "ssl_stream.h"

// The default config file name
#define DEFAULT_CONFIG_FILE_NAME "nettles.lua"


#ifdef HAVE_STRICMP
	#define os_stricmp(x,y) stricmp(x,y)
#else
	#ifdef HAVE__STRICMP
		#define os_stricmp(x,y) _stricmp(x,y)
	#else
		#ifdef HAVE_STRCASECMP
			#define os_stricmp(x,y) strcasecmp(x,y)
		#endif
	#endif
#endif

#define GLOBAL_DATA "globaldata"

#define LISTENER_TYPE_UNKNOWN 0
#define LISTENER_TYPE_CLEARTEXT 1
#define LISTENER_TYPE_CIPHERTEXT 2

// Functions that we expose to lua
int l_add_listener (lua_State *L);
int l_remove_listener (lua_State *L);

struct {
	char* name;
	lua_CFunction function;
} lua_function_table[] = {
		{"add_listener", l_add_listener},
		{"remove_listener", l_remove_listener},
		{NULL, NULL}
};


struct general_data {
	struct event_base *base;
	struct evdns_base *dns;
//	lua_State *L;
};

#ifdef TRACE
static void stackDump (lua_State *L) {
      int i;
      int top = lua_gettop(L);
      for (i = 1; i <= top; i++) {  /* repeat for each level */
        int t = lua_type(L, i);
        switch (t) {

          case LUA_TSTRING:  /* strings */
            printf("`%s'", lua_tostring(L, i));
            break;

          case LUA_TBOOLEAN:  /* booleans */
            printf(lua_toboolean(L, i) ? "true" : "false");
            break;

          case LUA_TNUMBER:  /* numbers */
            printf("%g", lua_tonumber(L, i));
            break;

          default:  /* other values */
            printf("%s", lua_typename(L, t));
            break;

        }
        printf("  ");  /* put a separator */
      }
      printf("\n");  /* end the listing */
    }
#endif

// Called by libevent when there is data to read
static void input_readcb(struct bufferevent *bev, void *user_data)
{
	stream_handle_t stream = user_data;

	stream_handle_t read = stream;
	stream_handle_t write = stream_get_other(stream);

	struct evbuffer* buffer = evbuffer_new();
	if (buffer == NULL)
	{
		// TODO Handle error
		fprintf(stderr, "Failed to allocate evbuffer");
		return;
	}

	error_type_t res = stream_read(read, buffer);
	if (res == ERROR_FAILED)
	{
		if (!stream_is_connected(read))
		{
			// Shutdown the other side
			stream_disconnect(write);

			// And free us
			stream_release_other(read);
			stream_free(read);
		}
	}
	else
	{
		if (evbuffer_get_length(buffer) > 0) {
			res = stream_write(write, buffer);
			if (res == ERROR_FAILED)
			{
				if (!stream_is_connected(write))
				{
					// Shutdown the other side
					stream_disconnect(read);

					// And free us
					stream_release_other(write);
					stream_free(write);
				}
			}
		}
	}

	evbuffer_free(buffer);
}

// Called by libevent when data has been sent
static void input_writecb(struct bufferevent *bev, void *user_data)
{
	stream_handle_t stream = user_data;

	stream_handle_t read = stream;
	stream_handle_t write = stream_get_other(stream);

	if (!stream_is_connected(read))
	{
		// Shutdown the other side
		stream_disconnect(write);

		// And free us
		stream_release_other(read);
		stream_free(read);
	}
}

// Called by libevent on various other events
static void input_eventcb(struct bufferevent *bev, short events, void *user_data)
{
	stream_handle_t stream = user_data;

	stream_handle_t read = stream;
	stream_handle_t write = stream_get_other(stream);

	if (events & BEV_EVENT_EOF)
	{
		stream_disconnect(read);
		printf("Connection closed.\n");
	}
	else if (events & BEV_EVENT_ERROR)
	{
		stream_disconnect(read);
		fprintf(stderr, "Connection error.\n");
	}
	else if (events & BEV_EVENT_CONNECTED)
	{
		stream_connected(read);
	}

	/* None of the other events can happen here, since we haven't enabled
	 * timeouts */

	if (!stream_is_connected(read))
	{
		// Shutdown the other side
		stream_disconnect(write);

		// And free us
		stream_release_other(read);
		stream_free(read);
	}
}

// Called by ssl_stream on successful authentication
static void authentication_cb(stream_handle_t stream, void* data)
{
	stream_handle_t other = stream_get_other(stream);
	struct bufferevent* bev = stream_get_bev(other);

	bufferevent_enable(bev, EV_WRITE | EV_READ);
}

// Create a stream
stream_handle_t create_stream(struct listener *data, int server)
{
	stream_handle_t stream = NULL;

	// Store stacksize so we can cleanup our mess...
	int stacksize = lua_gettop(data->L);

	lua_pushstring(data->L, "type");
    lua_gettable(data->L, -2);
    const char* type = lua_tostring(data->L, -1);
    lua_pop(data->L, 1);

    lua_pushstring(data->L, "certificate_file");
    lua_gettable(data->L, -2);
    const char* certificate_file = lua_tostring(data->L, -1);
    lua_pop(data->L, 1);

	lua_pushstring(data->L, "key_file");
    lua_gettable(data->L, -2);
    const char* key_file = lua_tostring(data->L, -1);
    lua_pop(data->L, 1);

	lua_pushstring(data->L, "trusted_ca_file");
    lua_gettable(data->L, -2);
    const char* trusted_ca_file = lua_tostring(data->L, -1);
    lua_pop(data->L, 1);

    int currentstacksize = lua_gettop(data->L);
    lua_pop(data->L, currentstacksize - stacksize);

    if (os_stricmp(type,"cleartext") == 0)
    {
    	// Cleartext
    	stream = stream_new_cleartext();
    	if (stream == NULL)
    	{
    		// Failed to allocate memory!
    		fprintf(stderr, "Failed to allocate memory!");
    		return NULL;
    	}

    }
    else if (os_stricmp(type,"ciphertext") == 0)
    {
    	// Ciphertext
    	if (server)
    	{
    		stream = stream_new_ssl_server();
    	}
    	else
    	{
    		stream = stream_new_ssl_client();
    	}
		if (stream == NULL)
		{
			// Failed to allocate memory!
			fprintf(stderr, "Failed to allocate memory!");
			return NULL;
		}

		stream_set_authentication_callback(stream, authentication_cb, stream);

		struct cipher_context* context = stream_get_cipher_context(stream);

	    int ret = x509parse_crtfile(&context->ca,trusted_ca_file);
	    if (ret != 0)
	    {
	    	// Fail
	    	stream_free(stream);
	    	return NULL;
	    }

	    ret = x509parse_crtfile(&context->certificate,certificate_file);
	    if (ret != 0)
	    {
	    	// Fail
	    	stream_free(stream);
	    	return NULL;
	    }

	    rsa_init( &context->rsa, RSA_PKCS_V15, 0 );
	    ret = x509parse_keyfile(&context->rsa,key_file,NULL);
	    if (ret != 0)
	    {
	    	// Failed
	    	stream_free(stream);
	    	return NULL;
	    }

	    ssl_set_own_cert(&context->ssl,&context->certificate,&context->rsa);
	    ssl_set_ca_chain(&context->ssl,&context->ca,&context->crl,NULL);

		if (server)
		{
			// TODO Read from file
			ssl_set_dh_param(&context->ssl, POLARSSL_DHM_RFC5114_MODP_1024_P, POLARSSL_DHM_RFC5114_MODP_1024_G);
		}

    }

	// Store stacksize so we can cleanup our mess...
	stacksize = lua_gettop(data->L);

	// Store the table so we can get it later
	lua_pushlightuserdata(data->L, stream);
	lua_pushvalue(data->L, -2);
	lua_settable(data->L, LUA_REGISTRYINDEX);

    currentstacksize = lua_gettop(data->L);
    lua_pop(data->L, currentstacksize - stacksize);

    return stream;
}

// Called by libevent when a client connects
static void listener_cb(struct evconnlistener *listener, evutil_socket_t fd,
    struct sockaddr *sa, int socklen, void *user_data)
{
	struct listener *data = user_data;

	assert(listener == data->listener);

	stream_handle_t read = NULL;
	stream_handle_t write = NULL;

#define BUFFERSIZE (100)
	char peer[BUFFERSIZE+1];
	memset(peer,0,BUFFERSIZE+1);

	// TODO handle ipv6
	if (sa->sa_family == AF_INET)
	{
		struct sockaddr_in* sin = (struct sockaddr_in*)sa;
		unsigned long addr = ntohl(sin->sin_addr.s_addr);
		char* addrptr = (char*)&addr;
		snprintf(peer,BUFFERSIZE,"%d.%d.%d.%d:%d",
				addrptr[3],
				addrptr[2],
				addrptr[1],
				addrptr[0],
				ntohs(sin->sin_port));
	}

	// Store stacksize so we can cleanup our mess...
	int stacksize = lua_gettop(data->L);

    lua_pushlightuserdata(data->L, (void *)data);  // push address
    lua_gettable(data->L, LUA_REGISTRYINDEX);  // retrieve value

    read = create_stream(data, 1);

	lua_pushstring(data->L, "on_accept");
    lua_gettable(data->L, -2);
	lua_pushstring(data->L, peer);

    if (lua_pcall(data->L, 1, 1, 0) != 0) {
		fprintf(stderr, "Failed to call on_accept! %s", lua_tostring(data->L, -1));
		goto error;
    }

    if (!lua_istable(data->L,-1))
    {
		fprintf(stderr, "Bad result from on_accept!");
		goto error;
    }

	lua_pushstring(data->L, "connect_ip");
    lua_gettable(data->L, -2);
    const char* connect_ip = lua_tostring(data->L, -1);
    lua_pop(data->L, 1);

	lua_pushstring(data->L, "connect_port");
    lua_gettable(data->L, -2);
    int connect_port = lua_tointeger(data->L, -1);
    lua_pop(data->L, 1);

    write = create_stream(data, 0);

    int currentstacksize = lua_gettop(data->L);
    lua_pop(data->L, currentstacksize - stacksize);

	// Link them together
	stream_add_other(read,write);
	stream_add_other(write,read);

	struct bufferevent* read_bev = bufferevent_socket_new(data->base, fd, BEV_OPT_CLOSE_ON_FREE);
	if (!read_bev) {
		fprintf(stderr, "Error constructing bufferevent!");
		goto error;
	}
	bufferevent_setcb(read_bev, input_readcb, input_writecb, input_eventcb, read);
//	bufferevent_enable(read_bev, EV_WRITE | EV_READ);
	bufferevent_disable(read_bev, EV_WRITE | EV_READ);

	stream_set_bev(read,read_bev);

	struct bufferevent* write_bev = bufferevent_socket_new(data->base, -1, BEV_OPT_CLOSE_ON_FREE);
	if (!write_bev) {
		fprintf(stderr, "Error constructing bufferevent!");
		goto error;
	}
	bufferevent_setcb(write_bev, input_readcb, input_writecb, input_eventcb, write);
//	bufferevent_enable(write_bev, EV_WRITE | EV_READ);
	bufferevent_disable(write_bev, EV_WRITE | EV_READ);

	stream_set_state(write, STATE_CONNECTING);
	int ret = bufferevent_socket_connect_hostname(write_bev,data->dns,AF_INET,connect_ip,connect_port);
	if (ret != 0)
	{
		goto error;
	}

	stream_set_bev(write,write_bev);

	// And tell the stream that we are connected
	stream_connected(read);

	return;

error:

	stream_release_other(read);
	stream_free(read);
	stream_release_other(write);
	stream_free(write);

	event_base_loopbreak(data->base);
	return;
}

// Add a new listener
void add_listener(const char* ip, int port, struct listener* listener)
{
	struct sockaddr_in sin;
	int sinlen = sizeof(struct sockaddr_in);

	memset(&sin, 0, sizeof(sin));

	int res = evutil_parse_sockaddr_port(ip, (struct sockaddr*)&sin, &sinlen);
	if (res != 0) {
		fprintf(stderr, "Could not parse address!\n");
		return;
	}
	sin.sin_port = htons(port);

	listener->listener = evconnlistener_new_bind(listener->base, listener_cb, listener,
	    LEV_OPT_REUSEABLE|LEV_OPT_CLOSE_ON_FREE, -1,
	    (struct sockaddr*)&sin,
	    sizeof(sin));

	if (!listener->listener) {
		fprintf(stderr, "Could not create a listener!\n");
		return;
	}

	return;
}

// Callback from lua for adding a listener
int l_add_listener (lua_State *L)
{
	if (!lua_istable(L, 1))
	{
		// TODO Fail
		return 0;
	}

	struct listener* listener = listener_new();
	if (listener == NULL)
	{
		return 0;
	}

	int stacksize = lua_gettop(L);

	lua_pushstring(L, GLOBAL_DATA);
	lua_gettable(L, LUA_REGISTRYINDEX);
	struct general_data* data = lua_touserdata(L, -1);

	listener->base = data->base;
	listener->dns = data->dns;
	listener->L = L;

	lua_pushstring(L, "listen_ip");
    lua_gettable(L, 1);
    const char* listen_ip = luaL_checkstring(L, -1);
	lua_pushstring(L, "listen_port");
    lua_gettable(L, 1);
    int listen_port = luaL_checkinteger(L, -1);

	add_listener(listen_ip, listen_port, listener);

	// Store table in registry
	lua_pushlightuserdata(L, listener);
	lua_pushvalue(L, 1);
	lua_settable(L, LUA_REGISTRYINDEX);

    int currentstacksize = lua_gettop(L);
    lua_pop(L, currentstacksize - stacksize);

    // Return the listener
    lua_pushlightuserdata(L, listener);

	return 1;  /* number of results */
}

// Callback from lua for removing a listener
int l_remove_listener (lua_State *L)
{
	if (!lua_islightuserdata(L, 1))
	{
		// TODO Fail
		return 0;
	}

	struct listener* listener = lua_touserdata(L,1);
	listener_free(listener);

	return 0;
}


int main(int argc, char* argv[])
{
	char* conf_filename = DEFAULT_CONFIG_FILE_NAME;
	if (argc > 1) {
		conf_filename = argv[1];
	}

	// Initialize libevent
#ifdef WIN32
	WSADATA wsa_data;
	WSAStartup(0x0201, &wsa_data);
#endif

	struct general_data* data = malloc(sizeof(struct general_data));
	if (data == NULL)
	{
		fprintf(stderr, "Failed to allocate memory for libevent!\n");
		return 1;
	}

	data->base = event_base_new();
	if (data->base == NULL)
	{
		fprintf(stderr, "Could not initialize libevent!\n");
		free(data);
		return 1;
	}

	data->dns = evdns_base_new(data->base, 1);
	if (data->dns == NULL)
	{
		fprintf(stderr, "Could not initialize libevent dns!\n");
		free(data);
		return 1;
	}

	// Initialize Lua
	lua_State *L = luaL_newstate();

	(void)luaopen_base(L);

	// Store table in registry
	lua_pushstring(L, GLOBAL_DATA);
	lua_pushlightuserdata(L, data);
	lua_settable(L, LUA_REGISTRYINDEX);

    // Add our own functions
	int i = 0;
	for (i = 0; lua_function_table[i].name != NULL;i++) {
		lua_pushcfunction(L, lua_function_table[i].function);
		lua_setglobal(L, lua_function_table[i].name);
	}

    // Parse config
    int ret = luaL_loadfile(L, conf_filename);
    if (ret != 0)
    {
    	return 1;
    }

    ret = lua_pcall(L, 0, LUA_MULTRET, 0);
    if (ret != 0)
    {
    	fprintf(stderr, "Failed to read conf: '%s'", lua_tostring(L, -1));
    	return -1;
    }

	(void)event_base_dispatch(data->base);

	// Close lua
    lua_close(L);

	// Close libevent
//	evconnlistener_free(ciphertext_listener);
//	evconnlistener_free(cleartext_listener);
	event_base_free(data->base);
	evdns_base_free(data->dns,0);

	free(data);

#ifdef WIN32
	WSACleanup();
#endif

	return EXIT_SUCCESS;
}
